#pragma once

#include <unordered_map>
#include <deque>

#include "Rdma.h"
#include "MemoryManager.h"
#include "Session.h"
#include "Cycles.h"


using std::unordered_map;
//flag

typedef void (*RpcResponseHandler)(void* self, void* ex, void* usr);
typedef void (*RpcRequestHandler)(void* self, void* ex, void* usr);
enum RpcFlag{
    RPC_SERVER = 0x1,
    RPC_RDMA = 0x1 << 1,
    RPC_SHARED_RECV_QUEUE = 0x1 << 2,
    RPC_PREALLOC = 0x1 << 3,
};

enum RpcError{
    RPC_ERROR_SUCCESS = 0,
    RPC_ERROR_TIME_OUT,
    RPC_ERROR_UNKNOWN_SESS_ID,
    RPC_ERROR_CLOSED_SESS,
    RPC_ERROR_LACK_REMOTE_MEMORY,
    RPC_LARGE_MSG_NOTIFY,
    RPC_LARGE_MSG_FINISH,
    RPC_LARGE_MSG_NOTIFY2,
    RPC_LARGE_MSG_FINISH2,
};

struct RpcMsg
{
    uint32_t msg_len; //msg_total_len
    uint32_t sync; //msg can be free?

    uint64_t sign; //signature
    uint32_t free;
    uint16_t rpc_type; //rpc request type
    uint16_t status_code; // rpc response status code
    
    uint32_t data_len;
    char data[0];
}__attribute__((packed));
static_assert(sizeof(RpcMsg) == 28, "   ");

struct RpcLargeMsgMeta
{
    uint8_t state;
    uint64_t trigger_request;
    uint64_t client_buffer_addr;
    //uint64_t client_buffer_len;
    uint64_t server_buffer_addr;
    //uint64_t server_buffer_len;
}__attribute__((packed));
static_assert(sizeof(RpcLargeMsgMeta) == 25, "  ");

struct RpcExecution {
    struct RpcMsg* req_msg;
    struct RpcMsg* resp_msg;

    RpcResponseHandler cb;
    void* cb_arg;
    int cb_error;

    uint32_t self_id;
 
    uint64_t trigger_req;
    
    uint64_t remote_addr;
    uint32_t large_msg_len;
    uint32_t finish_len;

    uint64_t time_stamp;
    void *wait_free_rpc;
    uint32_t wait_free_sess_id;
};


class MiniRpc
{
public:
    void* memory_; // null : upper-level not provide
    uint64_t memory_size_; // 0 : upper-level not provide
    uint8_t flag_;
    uint32_t max_msg_size_;

    MiniRpc() 
    : memory_(NULL), memory_size_(0), flag_(0), max_msg_size_(0){}
    ~MiniRpc(){}

};




class RdmaRpc : public MiniRpc
{ 
private:
    std::unordered_map<uint16_t, RpcRequestHandler> id2service_; // server registered function
    std::unordered_map<uint16_t, void *> id2serviceprovider_;
    uint32_t cur_sess_id_;
    std::unordered_map<uint32_t, Session*> id2sess_;
    int (*datapath_event_loop_)(); // event loop, used by rpc itself
    Session *server_listen_sess_; //socket to listen at server  
    int epoll_fd_; // return by epoll_create
    struct EpollEventCtx* events_; //pre-allocated memory used when epoll processing
    int total_inflight_;
    std::deque<RpcExecution*> large_write_msg_pending_;
    std::deque<RpcExecution*> large_read_msg_pending_;

    MemoryManager &mm_; // memory manager
    Rdma rdma_res_; //rdma context
   

    
public:
    RdmaRpc(MemoryManager &mm);
    ~RdmaRpc();
    void *AllocPiece() { return mm_.AllocFrag(); }
    int FreePiece(void* buf) { return mm_.FreeFrag(buf); }



    void* AllocMsg() {return mm_.AllocMsg(); }
    int FreeMsg(void* buf) { return mm_.FreeMsg(buf); }
    
    void* AllocLarge() {return mm_.AllocRing(); }
    int FreeLarge(void* buf) {return mm_.FreeRing(buf);}
    Session* CreateSess();
    uint32_t ConnectSess(const char* uri, RpcSessHandler cb, void* cb_arg);

    
    int CloseSess(uint32_t sess_id);


    int BindSess(const char* uri);
    int AcceptSess();

    int Init(uint8_t flag);
    void Destroy();

    


    int RpcServiceRegister(uint16_t rpc_type, RpcRequestHandler request_handler);
    int RpcServiceProviderRegister(uint16_t rpc_type, void *pro)
    {
        id2serviceprovider_[rpc_type] = pro;
        return 0;
    }
    int RpcAddRx(int num);
    int RpcServerHandleRx();
    int RpcServerHandleTx();
    void RpcServiceComplete(struct RpcExecution* ex);

    int RpcCliHandleRx();
    int RpcCliHandleTx();


    
    void* PrepareMsg(struct RpcMsg* header, uint16_t rpc_type);
    int RpcPollOnce();
    int RpcCliWriteReq(int sess_id, struct RpcMsg* req, RpcResponseHandler cb, void* cb_arg);
    int RpcBatchReq();
    void* get_rpc_response(struct RpcExecution* ex)
    {
        return static_cast<RpcMsg*>(ex->resp_msg)->data;

    }
    int free_rpc_request(struct RpcExecution* ex)
    {
        return FreeMsg(ex->req_msg);
    }
    bool is_successful(struct RpcExecution* ex)
    {
        return ex->cb_error == RPC_ERROR_SUCCESS;
    }

    Session* get_sess_ctx (uint32_t sess_id) 
    {
        auto iter = id2sess_.find(sess_id);
        if(iter == id2sess_.end())
        {
            return nullptr;
        }
        return iter->second;

    }
    int del_sess_ctx (uint32_t sess_id) 
    {
        auto iter = id2sess_.find(sess_id);
        if(iter == id2sess_.end())
        {
            return -1;
        }
        Session* sess = iter->second;
        sess->Destroy();
        delete sess;
        id2sess_.erase(sess_id);
        return 0;

    }
    void AllocRemoteLargeBuf(RpcMsg* msg)
    {
        msg->status_code = RPC_LARGE_MSG_NOTIFY2;

    }
private:
    int SendSmallMsg(Session *sess, RpcMsg *req); 
    int SendLargeMsg(Session *sess, RpcMsg *req);
    void WriteSegment();
    void ReadSegment();

};



